

import os

import torch
import json
import torchvision.transforms as transforms
from torch.autograd import Variable
import cv2 as cv
from tqdm import tqdm
from PIL import Image

from utils.inittestdataset import TestDatasetTools
from config.config import Config



def get_gt(json_file, id):

    with open (json_file, 'r') as gt_file:
        gt = json.load(gt_file)
        return gt[str(id)+'.wav']

def init_dataset(config):

    if os.path.exists(os.path.join(config.task2_test_path, '001')):
        print('Already Init.')
    else:
        print('Init Test Dataset')
        dataset_tools = TestDatasetTools(config.raw_task2_test_path, config.task2_test_path)
        dataset_tools.start_convert()


def vote():
    pass


def test(config):

    corr = 0

    model_1 = torch.load(os.path.join(config.model_root_path, config.task2_model_1))
    model_2 = torch.load(os.path.join(config.model_root_path, config.task2_model_2))

    img_transform = transforms.Compose(
                    [
                        transforms.Resize([224,224]),
                        transforms.ToTensor(),
                        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
                    ]
                )

    test_data_list = os.listdir(config.task2_test_path)
    # for test_data_idx in tqdm(range(len(test_data_list))):
    for test_data_idx in range(len(test_data_list)):
        first_flag = True
        test_data = test_data_list[test_data_idx]
        print('test data: {}'.format(test_data))
        the_gt = get_gt(config.task2_test_gt_path, test_data)
        print('gt: {}'.format(the_gt))
        # test_data: 001, 002, 003, ...
        test_dir = os.path.join(config.task2_test_path, test_data)
        test_img_list = os.listdir(test_dir)
        result = None
        for test_img_idx in range(len(test_img_list)):
            
            test_img_name = test_img_list[test_img_idx]
            test_img_path = os.path.join(test_dir, test_img_name)
            test_img_array = cv.imread(test_img_path)
            test_img = Image.fromarray(test_img_array)
            test_img_tensor = img_transform(test_img)
            test_img_tensor = Variable(torch.unsqueeze(test_img_tensor, dim=0).float(), requires_grad=False)

            if torch.cuda.is_available():
                test_img_tensor = test_img_tensor.cuda()
            
            # single_result_1 = model_1(test_img_tensor)
            # print(single_result_1)
            # _, pred = torch.max(single_result_1.data, 1)
            # print(pred+1)
            single_result_2 = model_2(test_img_tensor)
            # print(single_result_2)
            _, pred = torch.max(single_result_2.data, 1)
            # print(pred+1)
            if first_flag:
                result = single_result_2
                # result += single_result_2
                first_flag = not first_flag
            else:
                result += single_result_2
                # result += single_result_2
        # print(result)
        print('result: {}'.format(result.cpu().argmax()+1))
        if int(result.cpu().argmax()+1) == int(the_gt[2:]):
            print('Good')
            corr += 1
    print('Correct : {}/50'.format(corr))



if __name__ == '__main__':
    config = Config()
    init_dataset(config)
    test(config)
    